import cv2
import pickle
import matplotlib.pyplot as plt

# 加载保存的模型clf
with open("model.pkl", "rb") as file:
    clf = pickle.load(file)

# 读取原图和分割图
original_image = cv2.imread("Sandstone_2.tif", 0)
segment = cv2.imread("Sandstone_2_segment.tif", 0)

# 对原图进行预测
predicted_segment = clf.predict(original_image.reshape(-1, 1))
predicted_segment = predicted_segment.reshape(original_image.shape)

# 计算准确率
accuracy = (predicted_segment == segment).mean()
print("准确率:", accuracy)

# 显示原图、原分割图和推理分割图
fig, axes = plt.subplots(1, 3, figsize=(12, 4))

axes[0].imshow(original_image, cmap='gray')
axes[0].set_title("Original Image")

axes[1].imshow(segment, cmap='gray')
axes[1].set_title("Segment")

axes[2].imshow(predicted_segment, cmap='gray')
axes[2].set_title(f"Segmentation (acc: {accuracy:.3f})")

for ax in axes:
    ax.axis('off')

plt.tight_layout()
plt.show()
